"""
src/gdm.py
    Classes for Geom-Display-Mass (GDM) objects

Written By:
    James Thomas
    Email: jim@houseoftechnology.org
    Web: http://mission-cognition.houseoftechnology.org/

Copyright 2009-2012

This file is part of the PandaODElib distribution.

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public License
as published by the Free Software Foundation; either version 3 of
the License, or (at your option) any later version. The text of the
GNU Lesser General Public License is included with this library in
the file named LICENSE.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
"""
import os

from pandac.PandaModules import Quat, Vec3, Vec4
from pandac.PandaModules import OdeMass
from pandac.PandaModules import OdeGeom, OdeBoxGeom, OdeCylinderGeom, OdeSphereGeom, OdeTriMeshGeom
from pandac.PandaModules import OdeTriMeshData

import world
import odelib
import pvis

class Mesh(OdeTriMeshData):
    # Provide new so we can have subclasses that add parameters
    def __new__(cls, *args, **kwargs):
        # Initialize the TriMeshData class
        return OdeTriMeshData.__new__(cls)

    """A mesh object based on ode's TriMesh with added information for generating a display."""
    def __init__(self):
        OdeTriMeshData.__init__(self)
        self._verticies = []
        self._colors = []
        self._normals = []

    def build(self, Verticies, Faces, Colors=None, Normals=None, IsSoft=False):
        """Build a displayable tri-mesh object.  This is an overload of ode.TriMeshData.build()."""
        # Fill out our TriMeshData object properties by calling the parent
        OdeTriMeshData.build(self, Verticies, Faces)

        self._verticies = []
        self._colors = []
        self._normals = []

        # Expand the mesh to make it suitable for use with Visual
        for face in Faces:
            for index, vertex in enumerate(face):
                self._verticies.append(Verticies[vertex])
                if Colors is not None:
                    self._colors.append(Colors[vertex])
                else:
                    # Assign a default
                    self._colors.append(Vec4(1,1,1,1)) # White
                if Normals is not None:
                    self._normals.append(Normals[vertex])
                else:
                    # Compute the normals using cross products
                    v1 = Vec3(Verticies[face[index]])
                    v2 = Vec3(Verticies[face[(index + 1) % 3]])
                    normal = v1.cross(v2)
                    normal = normal if normal.normalize() else Vec3.unitX()
                    self._normals.append(normal)

        if IsSoft:
            # Need to go back and average normals -- implement later
            pass

    def MakeMeshDisplayObject(self):
        """Generate a VPython display object based on the tri-mesh properties."""
        return pvis.faces(pos=self._verticies, colors=self._colors, normals=self._normals)

class DisplayElement(object):
    """A Display-only element class -- no physics properties"""
    def __init__(self, DisplayObject=None):
        """Optionally set the panda3d model that will be associated with this display element."""
        self._disp = DisplayObject

    def SetVisible(self, IsVisible):
        """Set the visibility of the panda3d model. True=visible, False=not visible"""
        self._disp.visible = IsVisible

    def GetDisplayObject(self):
        """Get the panda3d model associated with this display element."""
        return self._disp

    def SetDisplayObject(self, Object):
        """Set the panda3d model that will be associated with this display element."""
        self._disp = Object
        return self

    def SetNode(self, Node):
        self._disp.node.wrtReparentTo(Node)
        
class GDMElement(object):
    """An element with display, physics properties, and a collision Geom."""
    def __init__(self):
        self._disp = None
        self._geom = None
        self._mass = OdeMass()

        self._path = '../resources' + os.sep
        
        # Setup a default system of no collisions -- we are using geoms to track
        # display object position and orientation so we are defining them but
        # collisions can blow up the sim so we have to make this safe by default.
        self._goemCategory = 0 # no category
        self._goemCollide  = 0 # collide with nothing
    
        self._space = world.World().space
        
        self._offsetPos = Vec3(0,0,0)

    def _translate(self, Offset):
        """Shift the geom (but not the mass) relative to the body/cg origin"""
        p1 = Vec3(self._geom.getOffsetPosition())
        self._geom.setOffsetPosition(self._offsetPos + Offset)

    def SetCollisionBitfields(self, CategoryBits, CollideBits):
        """Set the category and collide bitfields for this GDM. 
        Each GDM has a "category" and "collide" bitfield that can be used to assist
        the space algorithms in determining which GDMs should interact and which
        should not. Use of this function is optional, by default GDMs are configured
        for no collisions (category is bit 31, collide with bits 0-30). The test for
        a collision is ((cat1 & col2) || (cat2 & col1)).
        Each bit position in the bitfield represents a different category of object.
        The actual meaning of these categories (if any) is user defined. The category
        bitfield indicates which categories a GDM is a member of. The collide bitfield
        indicates which categories the GDM will collide with during collision detection."""
         
        self._goemCategory = CategoryBits
        self._goemCollide  = CollideBits
        
        # Update the settings on our geom if it has been defined already.
        if self._geom is not None:
            self._geom.setCategoryBits(self._goemCategory)
            self._geom.setCollideBits(self._goemCollide)

    def SetVisible(self, IsVisible):
        """Set the visibility of the panda3d model. True=visible, False=not visible"""
        self._disp.visible = IsVisible

    # The goem holds the reference position
    def setOffsetPosition(self, Position):
        """Set the offset (to the body) position of the GDM in local coordinates."""
        self._geom.setOffsetPosition(Position)
        self._mass.translate(Position)
        self._offsetPos = Vec3(Position)
        
    # The geom holds the reference orientation
    def setOffsetQuaternion(self, Quat):
        """Set the offset (to the body) orientation of the GDM in local coordinates."""
        self._geom.setOffsetQuaternion(Quat)
        # TODO: Need to do something to the mass
        # self._mass.rotate()
        self._offsetQuat = Quat
        
    def GetDisplayObject(self):
        """Get the Panda3d display object associated with this display element."""
        return self._disp

    def GetGeom(self):
        """Get the geom assocaiated with this element."""
        return self._geom
    
    def GetMass(self):
        """Get the mass of this element."""
        return self._mass
    
    def DefineCustom(self, Geom, Mass, Disp):
        assert isinstance(Geom, OdeGeom)
        assert isinstance(Mass, OdeMass)
        assert isinstance(Disp, pvis.Shape)
        
        self._geom = Geom
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)
        self._mass = Mass
        self._disp = Disp
        return self
    
    def DefineBox(self, Density, SizeX, SizeY, SizeZ, CustomDisp=None):
        """Define this element as a ode Box."""
        self._geom = OdeBoxGeom(self._space, (SizeX, SizeY, SizeZ)) # BigSpace was None, problem with local body spaces
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)
        
        if CustomDisp is None:
            self._disp = pvis.box(length=SizeX, height=SizeY, width=SizeZ)
        elif isinstance(CustomDisp, pvis.Shape):
            self._disp = CustomDisp
        elif issubclass(CustomDisp, pvis.Shape):
            self._disp = CustomDisp(length=SizeX, height=SizeY, width=SizeZ)
        else:
            raise TypeError, 'CustomDisp %s is not an instance or sublcass of pvis.Shape' % type(CustomDisp)
              
        self._mass = OdeMass()
        self._mass.setBox(Density, SizeX, SizeY, SizeZ)
        return self

    def DefineBoxTotal(self, TotalMass, SizeX, SizeY, SizeZ, CustomDisp=None):
        """Define this element as an ode BoxTotal."""
        self._geom = OdeBoxGeom(self._space, (SizeX, SizeY, SizeZ))
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)

        if CustomDisp is None:
            self._disp = pvis.box(length=SizeX, height=SizeY, width=SizeZ)
        elif isinstance(CustomDisp, pvis.Shape):
            self._disp = CustomDisp
        elif issubclass(CustomDisp, pvis.Shape):
            self._disp = CustomDisp(length=SizeX, height=SizeY, width=SizeZ)
        else:
            raise TypeError, 'CustomDisp %s is not an instance or sublcass of pvis.Shape' % type(CustomDisp)
        
        self._mass = OdeMass()
        self._mass.setBoxTotal(TotalMass, SizeX, SizeY, SizeZ)
        return self

    def DefineSphere(self, Density, Radius, CustomDisp=None):
        """Define this element as an ode Sphere."""
        self._geom = OdeSphereGeom(self._space, Radius)
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)

        if CustomDisp is None:
            self._disp = pvis.sphere(radius=Radius)
        elif isinstance(CustomDisp, pvis.Shape):
            self._disp = CustomDisp
        elif issubclass(CustomDisp, pvis.Shape):
            self._disp = CustomDisp(radius=Radius)
        else:
            raise TypeError, 'CustomDisp %s is not an instance or sublcass of pvis.Shape' % type(CustomDisp)

        self._mass = OdeMass()
        self._mass.setSphere(Density, Radius)
        return self

    def DefineSphereTotal(self, TotalMass, Radius, CustomDisp=None):
        """Define this element as an ode SphereTotal."""
        self._geom = OdeSphereGeom(self._space, Radius)
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)

        if CustomDisp is None:
            self._disp = pvis.sphere(radius=Radius)
        elif isinstance(CustomDisp, pvis.Shape):
            self._disp = CustomDisp
        elif issubclass(CustomDisp, pvis.Shape):
            self._disp = CustomDisp(radius=Radius)
        else:
            raise TypeError, 'CustomDisp %s is not an instance or sublcass of pvis.Shape' % type(CustomDisp)

        self._mass = OdeMass()
        self._mass.setSphereTotal(TotalMass, Radius)
        return self

    def DefineCylinder(self, Density, Radius, Length, CustomDisp=None):
        """Define this element as an ode Cylinder."""
        self._geom = OdeCylinderGeom(self._space, Radius, Length)
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)

        if CustomDisp is None:
            self._disp = pvis.cylinder(radius=Radius, length=Length)
        elif isinstance(CustomDisp, pvis.Shape):
            self._disp = CustomDisp
        elif issubclass(CustomDisp, pvis.Shape):
            self._disp = CustomDisp(radius=Radius, length=Length)
        else:
            raise TypeError, 'CustomDisp %s is not an instance or sublcass of pvis.Shape' % type(CustomDisp)

        self._mass = OdeMass()
        self._mass.setCylinder(Density, 3, Radius, Length) # 3 is z-axis -- same as geom
        return self

    def DefineCylinderTotal(self, TotalMass, Radius, Length, CustomDisp=None):
        """Define this element as an ode CylinderTotal."""
        self._geom = OdeCylinderGeom(self._space, Radius, Length)
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)

        if CustomDisp is None:
            self._disp = pvis.cylinder(radius=Radius, length=Length)
        elif isinstance(CustomDisp, pvis.Shape):
            self._disp = CustomDisp
        elif issubclass(CustomDisp, pvis.Shape):
            self._disp = CustomDisp(radius=Radius, length=Length)
        else:
            raise TypeError, 'CustomDisp %s is not an instance or sublcass of pvis.Shape' % type(CustomDisp)
        
        self._mass = OdeMass()
        self._mass.setCylinderTotal(TotalMass, 3, Radius, Length) # 3 is z-axis -- same as geom
        return self

    def DefineMeshTotal(self, TotalMass, CenterOfMass, InertiaMatrix, Mesh):
        """Define this element as a mesh.  The Mesh must be a ode3d.Mesh class object.  The
        cg and IntertiaMaxtrix are defined in world coordinates."""
        self._geom = OdeTriMeshGeom(Mesh, self._space)
        self._geom.setCategoryBits(self._goemCategory)
        self._geom.setCollideBits(self._goemCollide)

        self._disp = Mesh.MakeMeshObject()

        # Setup the mass properties
        self._mass = OdeMass()
        cgx, cgy, cgz = CenterOfMass
        (I11, I12, I13), (I21, I22, I23), (I31, I32, I33) = InertiaMatrix
        self._mass.setParameters(TotalMass, cgx, cgy, cgz, I11, I22, I33, I12, I13, I23)
        return self

    def UpdateDisplay(self):
        # GDMEs are not put into the body display frames.  Thus the
        # display object is being set in absolute global coords
        self._disp.q = self._geom.getQuaternion()
        self._disp.pos = self._geom.getPosition()
        
        #self._disp.setQuat(self._geom.getOffsetQuaternion() * self._geom.getQuaternion())
        #p1 = self._geom.getBody().vectorToWorld(self._geom.getOffsetPosition())
        #self._disp.setPos(p1 + self._geom.getPosition())
        #print self._geom.getPosition(), self._geom.getBody().getPosition()
    
    def dump(self):
        print 'id=%d, class=%s' % (id(self), self.__class__.__name__)
        print self._disp
        print self._disp.q
        print self._disp.pos
        print self._geom.getPosition()
        print self._geom.getQuaternion()
        print
        
class GDMBody(odelib.BaseBody):
    """A class which defines an ode-like body which has one or more
    GDM elements and display elements associated with it."""
    def __init__(self, World):
        # Initialize the Body class
        odelib.BaseBody.__init__(self, World)

        self._elementDict = {}

        # Mass not defined yet
        self._myNetMass = None

        # My display node (used to group my GDMEs)
        self._myNode = World.scene.attachNewNode("GDMBodyContainer")
        
        # Keep track of our body cg to user origin offset
        # ODE has a major limitation in that the cg must be at the 
        # body origin.  This is extremely confusing when building a
        # GDMBody using multiple GDMs which may shift the cg.  Thus
        # we will keep track of the cg offset and shift all of the
        # geoms relative to the body in order to move the cg back to
        # the body origin.  Hopefully the user won't notice.
        self._cgOffset = Vec3(0,0,0)

    def _recalcMass(self):
        # Recalculate the mass as the user constructed it
        netMass = OdeMass()
        for element in self._elementDict.itervalues():
            if isinstance(element, GDMElement):
                netMass.add(element.GetMass())

        # if the cg is not on the origin of the body we need to apply an
        # offset so that they align.  
        self._cgOffset = -Vec3(netMass.getCenter())
        netMass.translate(self._cgOffset)

        # Assign the mass to our body.  Use the base class call
        # since we are blocking use of the function at this class level
        odelib.BaseBody.setMass(self, netMass)
        
        # Now we need to offset all the geoms to compensate for the
        # difference in coordinate systems
        for element in self._elementDict.itervalues():
            if isinstance(element, GDMElement):
                element._translate(self._cgOffset)

        
    # Overload the setPosition function so we can make sure that the display frame matches
    def setPosition(self, Position):
        """Set the position of the body in world coordinates (the display object is moved to match)."""
        odelib.BaseBody.setPosition(self, Position)
        self._myNode.setPos(*Position)

    # Overload the setQuaternion function so we can make sure that the display frame matches
    def setQuaternion(self, q):
        """Set the orientation of the body in world coordinates (the display object is oriented to match)."""
        odelib.BaseBody.setQuaternion(self, q)
        self._myNode.setQuat(q)

    def AddElement(self, ElementKey, Element):
        """Add a GDM/display element to this body and use ElementKey as a handle to it."""
        if ElementKey in self._elementDict:  # Not ready for replaces yet
            raise TypeError, 'Element %s already exists in this Body' % ElementKey

        # Type check to avoid insidious bugs        
        if not (isinstance(Element, GDMElement) | isinstance(Element, DisplayElement)):
            raise TypeError, 'Must be an instance of GDMElement or DisplayElement'

        # Add to our element dictionary
        self._elementDict[ElementKey] = Element

        if isinstance(Element, DisplayElement):
            # Put this element into our display node
            Element.SetNode(self._myNode)
            
        elif isinstance(Element, GDMElement):
            Element.GetGeom().setBody(self)
            self._recalcMass()
            # Put this into the scene (since they use global coordinates)
            #Element.SetNode()

    def GetElementKeys(self):
        """Returns a list of element keys associated with this body."""
        return self._elementDict.keys()

    def SetElementVisible(self, ElementKey, IsVisible):
        """Set the visibility of the specified Element.  True=visible, False=not visible, None=toggle."""
        assert ElementKey in self._elementDict
        self._elementDict.get(ElementKey).SetVisible(IsVisible)

    def RemoveElement(self, ElementKey):
        """Remove (and destroy) the specified element from the body."""
        assert ElementKey in self._elementDict
        element = self._elementDict.pop(ElementKey)

        # Get rid of the display portion
        element.GetDisplayObject().removeNode()

        # Unlink the Geom from its body disable it and remove it from it's space
        element.GetGeom().setBody(None)
        element.GetGeom().disable()
        self._mySpace.remove(element.GetGeom())

        # Recompute the mass using the remaining features
        self._myNetMass.setZero()
        for element in self._elementDict.itervalues():
            self._myNetMass.add(element.GetMass())
        self.setMass(self._myNetMass)

    def GetFeature(self, ElementKey=''):
        """Returns the display object associated with the specified element key.  Returns
        the body display node object if the key is blank."""
        if ElementKey in self._elementDict:
            return self._elementDict.get(ElementKey).GetDisplayObject()
        else:
            return self._myNode

    def GetSpace(self):
        """Returns the ode collision space associated with this body."""
        return self._mySpace

    def SetVisible(self, IsVisible):
        """Sets the visibility for every element associated with the body.
        True=visible, False=not visible."""
        #for key in self._elementDict:
        #    self.SetElementVisible(key, IsVisible)

        if IsVisible:
            self._myNode.show()
        else:
            self._myNode.hide()

    def UpdateDisplay(self):
        """Sync's the VPython display with the bodies ode physics state."""
        # If we reset position and quaternion it will update
        # the display frame which has our display elements in it.
        self.setPosition(self.getPosition())
        self.setQuaternion(Quat(self.getQuaternion())) # Odd that these are not directly compatible
        
        # Call update display for each of our GDM elements
        for element in self._elementDict.itervalues():
            if isinstance(element, GDMElement):
                element.UpdateDisplay()

    # Functions that should not be used anymore
    def setMass(self, M):
        """An override to prevent use of this function.  Throws an assertion if called.
        Assign GDMElements to the body to for it to have mass."""
        assert False # Do not use this function anymore

    def setRotation(self, R):
        """An override to prevent use of this function.  Throws an assertion if called.
        This could be implemented in the future if somebody wants to figure out how to do it.
        You have to use setQuaternion() instead -- sorry."""
        assert False # Do not use this function anymore

    def dump(self):
        for name, element in self._elementDict.iteritems():
            if isinstance(element, GDMElement):
                print 'Element: ', name
                element.dump()
                print
        
class GDMAssembly(odelib.Assembly):
    """An assembly of BaseBody, GDMBody, Assembly, GDMAssembly objects used for complex structures.
    It is expected that all the bodies/sub-assemblies in the assembly will be connected
    to each other by some sort of joint (managed by the user of this class)."""
    def __init__(self):
        # Initialize the odelib.Assembly class
        odelib.Assembly.__init__(self)

        self._featureMap = {} # a map of display object to body key

        # My display node (used to group my bodies / sub-assemblies displays)
        self._myNode =  world.World().scene.attachNewNode("GDMAssemblyContainer")

    def AddObj(self, Name, Obj):
        """Add a obj to the assembly using Name for future reference.  Name must be a string
        and must not include a period because you may not use dotted sub-assembly keys here."""
        if not isinstance(Name, str) or (Name.find('.') >= 0):
            raise TypeError, 'Name must be a string and must not include a period because you may not use dotted sub-assembly keys here.'
        
        odelib.Assembly.AddObj(self, Name, Obj)

        # Objects derived from BaseBody or Assembly don't have features or display nodes
        if hasattr(Obj, 'GetFeature'):
            # Add the obj to our lookup by pvis feature map
            self._featureMap[Obj.GetFeature()] = Name

        if hasattr(Obj, '_myNode'):
            # Reparent the display node to ourselves
            Obj._myNode.wrtReparentTo(self._myNode)

    def SetVisible(self, Name, isVisible):
        """Set the visibility of the Body/Assembly associated with Name.  
        Name supports the dotted sub-assembly format."""
        obj = self.GetObj(Name)

        # Objects derived from BaseBody or Assembly don't have displays
        if hasattr(obj, 'SetVisible'):
            obj.SetVisible(isVisible)

    def GetFeature(self, ElementKey=''):
        """Return the display element feature for the named object.  If ElementKey is blank then
        the display node associated with this assembly is returned.  Use the assembly/subassembly
        dotted notation to select which object and display element are selected"""
        
        # Handle dotted keys
        splits = ElementKey.split('.', 1)
        
        if len(splits) == 1:
            splits.append('')
        
        key, remainder = splits
        
        if key:
            return self.GetObj(key).GetFeature(remainder)
        else:
            return self._myNode

    def Pick(self, PickObject):
        """Determine which body the display object is associated with."""

        feature = self._featureMap.get(PickObject)

        if feature is not None:
            return feature
        
        # Call Pick on each of our objects
        for obj in self._objDict.itervalues():
            if hasattr(obj, 'Pick'):
                feature = obj.Pick(PickObject)
        
            if feature is not None:
                return feature
            
        return None

    def UpdateDisplay(self):
        """Make the display for all of the bodies in this assembly consistent with the ode physics.""" 
        for obj in self._objDict.itervalues():
            try:
                # Update the body Display
                obj.UpdateDisplay()
            except AttributeError:
                pass # Objects derived from BaseBody don't have an UpdateDisplay function

    def dump(self):
        for name, obj in self._objDict.iteritems():
            try:
                print 'Obj: ', name
                obj.dump()
                print
            except AttributeError:
                pass
